#! /usr/bin/env python
# coding=utf-8


# 读取target_lista 
# 确保一个科 只有一种 长度为最长
# 生成两个输出的fasta文件 同时 
import string
from Bio import SeqIO
from Bio import Seq
import sys
import argparse



dis = '''
用于对提取出来的mRNA 筛选特定物种 保证每个科 只有一个长度最长的代表性序列 同时修改缩写名称 加入 目 和 科 为 目两位 科三位
tblastn_fetch_gene01_one_seq_one_family.py -p filter -s spe_table_add1.tsv -i raw_mRNA.fa
'''


parser = argparse.ArgumentParser(description=dis)





parser.add_argument('-p',
                help='必须给定，结果前缀')

parser.add_argument('-s',
                help='必须给定，物种列表 用于提取 ')



parser.add_argument('-i',
                help='必须给定，需要输入的mRNA文件')

# 获取参数
args = parser.parse_args()



if not args.p or not args.s or not args.i:
    parser.print_help()
    sys.exit()




family_order = []

spe_table = args.s#'spe_table'

mRNA_file = args.i#r'raw_mRNA.fa'


prefix = args.p

out_mRNA=  open(prefix + r'_mRNA.fa','w')
out_pep  = open(prefix + r'_pep.fa','w')


abb_dic = {}
abb_order = []
family_abb_dic = {}


with open(spe_table) as fila:
	for i in fila:
		k = i.strip().split('\t')
		if len(k)>2:
			abb = k[3]
			family = k[1]
			
			if family not in family_order:
				family_order.append(family)
			spe = k[0]
			order = k[2]
			new_abb = string.capwords(order[:2])+string.capwords(family[:3])+abb
			abb_dic[abb] = {'family':family,'order':order,'spe':spe,'new_abb':new_abb}

seq_dic=  {}


for i in SeqIO.parse(mRNA_file,'fasta'):
	abb = str(i.name).split('_')[0]
	if abb in abb_dic:
		seq_dic[abb] = str(i.seq)

for abb in seq_dic:
	
	family = abb_dic[abb]['family']
	if family not in family_abb_dic:
		family_abb_dic[family] = abb
	else:
		if len(seq_dic[family_abb_dic[family]]) < len(seq_dic[abb]):
			family_abb_dic[family] = abb

# 以fasta输出为一个临时文件
for i in family_order:
	#print(i,family_abb_dic[i])
	if i in family_abb_dic:
		abb = family_abb_dic[i]
		new_abb = abb_dic[abb]['new_abb']
		mRNA_seq = seq_dic[abb]
		pep_seq = str(Seq.Seq(seq_dic[abb]).translate())
		out_pep.write('>'+new_abb+'\n'+pep_seq+'\n')
		out_mRNA.write('>'+new_abb+'\n'+mRNA_seq+'\n')




out_pep.close()
out_mRNA.close()